import pandas as pd
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.feature as cfeature
import shapely.geometry as sgeom
from datetime import datetime,timedelta
import seaborn as sns
from global_land_mask import globe
from gsj_typhoon import tydat,see,count_rapidgrow,split_str_id,tydat_CMA,load_land_polygons,detect_landfall
import matplotlib.ticker as ticker
from tqdm import tqdm
from typlot.config.global_config import *


#operation config
# names = ['mojie_28','dusurui_16','gaemi_09','haikui_38','kangni_54','shantuo_44','saola_25','koinu_49']
# tynames,tyids = split_str_id(names)
# all_ini_time_mode = ['00_12','12','00']
# all_obs_baseline = ['RI','land']   # RI  land
# draw_obs_opt = True
# draw_opt = False



# debug config
names = ['mojie_28','dusurui_16','gaemi_09','haikui_38','kangni_54','shantuo_44','saola_25','koinu_49']
tynames,tyids = split_str_id(names)
all_ini_time_mode = ['00_12']
all_obs_baseline = ['RI'] 
draw_obs_opt = True
draw_opt = False

            


def add_red_dots(obs_arr: np.ndarray , ini_time_mode : str ) :
    ''' 补全观测红点，策略为先y后x循环
    obs_arr: 拓展到正确形状之后的数组。（指最上条红线恰好过原点）
    '''
    assert ini_time_mode in ['00','12','00_12'],'Wrong ini_time_mode value'
    assert obs_arr.ndim==2, 'Wrong obs_arr shape, ndim should be 2'
    ny,nx = obs_arr.shape
    obs_arr = obs_arr[::-1]  # 代码逻辑更直观一些
    dx = 2 if ini_time_mode=='00_12' else 1
    for j in range(ny-1,-1,-1):
        for i in range(nx):
            try:
                if obs_arr[j+1,i+dx] ==1 :
                    obs_arr[j,i]=1
            except: 
                pass #懒得处理数组越界问题了
    return obs_arr[::-1]



for ini_time_mode in all_ini_time_mode:
    for obs_baseline in all_obs_baseline:
        print('Setting:',ini_time_mode,obs_baseline)
        draw_data_list = []
        obs_list = []
        
        # 循环处理单台风，并将每个台风结果存储到列表中，用于后期统一处理 。
        for tyname,tyid  in zip(tynames,tyids):
            '''settings'''
            dir_dsr = os.path.join( global_ensdir,f"{tyname}_{tyid}")
            name_date=sorted(os.listdir(dir_dsr))
            name_date = [i for i in name_date if i[-2:] in ini_time_mode]
            dir_date= [os.path.join(dir_dsr,f) for f in name_date ]
            RIstd = 7  # std = 15   
            member_num = 52
            obs_path = os.path.join(global_obsdir,f'{tyname}_CMAobs.txt')
            pic_savepath = os.path.join(global_picdir,f"{obs_baseline}_heatmap",f"{ini_time_mode}","81")
            os.makedirs(pic_savepath,exist_ok=True)

            # 读取所有集合预报成员
            num_array = np.zeros((len(name_date),member_num),dtype=int)
            num=[]
            for date in dir_date:
                name_all = [f for f in os.listdir(date) if f.startswith("TRACK")]
                sorted_names = sorted(name_all, key=lambda x: int(x.split('TRACK_ID_')[-1]))
                path_all = [os.path.join(date, f) for f in sorted_names ]
                num2=[]
                for path in path_all:
                    t = tydat(path,RIstd)
                    nn = int(path.split("TRACK_ID_")[-1])
                    num2.append([t,nn]) 
                num.append(num2)
            
            # 遍历所有集合预报成员，确定最长leadtime， 用于确定数组shape
            a=[];b=[]
            for s in range(0,len(name_date)):
                start_time=datetime.strptime(name_date[s],"%Y%m%d%H")
                for i in range(0,len(num[s])):
                    a.append( max(num[s][i][0].time ))
                    b.append( max(num[s][i][0].time) - start_time )
            max_enddate = max(a) 
            max_dayrange = max(b)
                # draw_data( len_initime, len_lead_time )
            draw_data = np.zeros( (len(name_date),max_dayrange.days+1) ,dtype=int)
            
            
            
            
            # 填充单台风数组   draw_data_obs
            if draw_obs_opt==True:
                obs = tydat_CMA(obs_path) 
                    # 观测也用预报一样的数组，方便画图
                draw_data_obs = np.zeros( (len(name_date),max_dayrange.days+1),dtype=int )
                tt,dd = draw_data_obs.shape
                if obs_baseline=='RI':
                    record_obs = count_rapidgrow(RIstd,obs.umax, obs.time)
                    for t in range(tt):
                        for d in range(dd):
                            start_time = datetime.strptime(name_date[t],"%Y%m%d%H")
                            cal_time = start_time + d*timedelta(days=1)
                            calend_time = cal_time + timedelta(days=1)
                            trange = (obs.time>=cal_time) & (obs.time<calend_time)
                            draw_data_obs[t][d] = 1 if np.sum(record_obs[trange])>0 else 0 
 
                if obs_baseline=='land':
                    if obs_baseline == 'land':
                        # get  land  time
                        land_polys = load_land_polygons(os.path.join(global_shpdir,'China','bou1_4p.shp'), os.path.join(global_shpdir,'China','bou1_4p.dbf'))
                        flags,id1 = detect_landfall(obs.lat, obs.lon, land_polys)
                        if id1 is None:
                            pass
                        else:
                            land_time = obs.time[id1]
                            # identify  position
                            for t in range(tt):
                                for d in range(dd):
                                    start_time = datetime.strptime(name_date[t],"%Y%m%d%H")
                                    cal_time = start_time + d*timedelta(days=1)
                                    calend_time = cal_time + timedelta(days=1)
                                    trange = (obs.time>=cal_time) & (obs.time<calend_time)
                                    # 分类trange； 即有无登陆点
                                    draw_data_obs[t][d] = 1 if land_time in obs.time[trange] else 0
                                    
            # 收集单台风的观测, 15是指的15天leadtime
            obs_list.append(draw_data_obs.T[:15])
            
            #  填充单台风集合预报结果
            tt,dd = draw_data.shape 
            for t in tqdm(range(tt),desc=f'processing {tyname}' ):             #起报
                for d in range(dd):                #预报时效
                    start_time = datetime.strptime(name_date[t],"%Y%m%d%H")
                    cal_time = start_time + d*timedelta(days=1)
                    calend_time = cal_time + timedelta(hours=18)
                    n=0                          #成员
                    for i in range(len(num[t])): #对成员数循环
                        data = num[t][i][0]
                        RI = data.num_rapidgrow()
                        dayRI = RI[ (data.time>=cal_time) & (data.time<=calend_time) ]
                        if np.sum(dayRI>=1) :
                            n+=1
                    draw_data[t][d]=n
            
            # 收集单台风集合预报结果，15是 leadtime
            draw_data_list.append( draw_data.T[:15] )

        
        
        
        # 处理含有8个台风数据的列表至形状一致
        mode = 'upall'
        if mode == 'ori':
            max_dimx =max([arr.shape[1] for arr in draw_data_list])
            full_draw_data_list = [np.pad(arr,((0,0),(0,max_dimx-arr.shape[1])),mode='constant',constant_values=0)   for arr in draw_data_list]
            draw_data = np.array(full_draw_data_list).sum(axis=0)
            
            ###  draw num of RI points
            plt.figure(figsize=(15, 6))
            ax = sns.heatmap(draw_data,annot=True,
                             annot_kws={'size': 6},fmt="d", cmap="YlGnBu", 
                             cbar_kws={'label': 'number of rapid intensification'}, 
                             linewidths=0.5)  #,vmin=0,vmax =30
            
            ax.set_ylabel("Forecast realative to start time", fontsize=12)
            ax.set_xlabel("Start time from init time: h ", fontsize=12)
            #  手动设置xticks数量,snsheatmap default上限是19个
            ax.set_xticklabels(np.arange(draw_data.shape[1])*12) 
            plt.yticks(rotation=0)
            ylabels=[str(i)+"d"  for i in range(1,16)]
            ax.set_yticklabels(ylabels)
            ax.set_title('sum RI numbers of all members for 8 typhoons east asia')
            cbar = ax.collections[0].colorbar
            cbar.locator = ticker.MaxNLocator(integer=True)
            cbar.update_ticks()
            
            ### draw obs
            if draw_obs_opt == True : 
                for draw_data_obs in obs_list:
                    points = np.where(draw_data_obs == 1)
                    y_coords = points[0]+0.5
                    x_coords = points[1]+0.15
                    ax.scatter(x_coords, y_coords, color='red', marker='o',s=10)
            
            # plt.savefig(f"{pic_savepath}\\heatmap_num_days_startime",dpi=900)
            plt.show()
            # plt.close()
        
        
        # 处理 ens obs;  ens.shape = [len(ty)] (len(lead),len(init)) ; obs.shape= [len(ty)] ( (len(lead),len(init) )
        if mode=='upall':
            assert len(draw_data_list)==len(obs_list), 'obs and ensemble not corespond'
            new_obs_list,new_mem_list = [],[]
            # 处理8个台风的观测情况
            for n in range(len(draw_data_list)):  # tys
                obs = obs_list[n]
                mem = draw_data_list[n]
                # 该台风无红点
                if sum(sum(obs)) == 0:
                    new_mem_list.append(mem)
                    new_obs_list.append(obs)
                # 红点在左边界
                elif sum(obs[:,0])>0:
                    topindex = np.min(np.where(obs[:,0]==1))
                    # 左边界类型1, 水平单点
                    if obs[topindex,1]==0:
                        day_ahead = obs.shape[0]-1-topindex
                        x_compensate = 1+day_ahead*2 if ini_time_mode=='00_12' else day_ahead
                        padded_obs = np.pad(obs, pad_width=((0, 0), (x_compensate, 0)),mode='constant', constant_values=0)
                        padded_mem = np.pad(mem, pad_width=((0, 0), (x_compensate, 0)),mode='constant', constant_values=0)
                        new_obs_list.append(add_red_dots(padded_obs,ini_time_mode))
                        new_mem_list.append(padded_mem)
                    # 左边界类型2, 水平双点
                    elif obs[topindex,1]==1:
                        day_ahead = obs.shape[0]-1-topindex
                        x_compensate = day_ahead*2 if ini_time_mode=='00_12' else day_ahead
                        padded_obs = np.pad(obs, pad_width=((0, 0), (x_compensate, 0)),mode='constant', constant_values=0)
                        padded_mem = np.pad(mem, pad_width=((0, 0), (x_compensate, 0)),mode='constant', constant_values=0)
                        new_obs_list.append(add_red_dots(padded_obs,ini_time_mode))
                        new_mem_list.append(padded_mem)
                # 左边界没有红点,但是下边界存在红点
                elif sum(obs[:,0])==0 and sum(sum(obs)) != 0:
                    leftindex = np.min(np.where(obs[-1,:]==1))
                    new_mem_list.append(mem[:,leftindex:])
                    new_obs_list.append(obs[:,leftindex:])
                else: 
                    raise ValueError('Wrong here')
                
            
            max_xlength = max([mem.shape[1] for mem in new_mem_list])
            new_mem_list = [np.pad(mem, pad_width=((0, 0), (0, max_xlength-mem.shape[1])),mode='constant', constant_values=0) 
                            for mem in new_mem_list]
            # for typhoon koinu:    max_xlength-19    mem.shape[1]-20
            new_obs_list = [np.pad(obs, pad_width=((0, 0), (0, max_xlength-obs.shape[1])),mode='constant', constant_values=0) 
                            for obs in new_obs_list]
            
            #  draw heatmap of RI number
            draw_data = np.array(new_mem_list).sum(axis=0)
            
            
            
                        
        ''' draw '''
        draw_data = np.array(new_mem_list).sum(axis=0)
        # print(f"裁剪前数据形状: {draw_data.shape}")
        
        
        # 🎯 步骤1：先找到最大的maxindex
        maxindex = 0
        if draw_obs_opt == True and len(new_obs_list) > 0:
            for draw_data_obs in new_obs_list:
                points = np.where(draw_data_obs == 1)
                if len(points[1]) > 0:  # 确保有观测点
                    current_max = points[1].max()
                    maxindex = max(maxindex, current_max)
        
        # print(f"找到的maxindex: {maxindex}")
        
        # 🎯 步骤2：根据maxindex裁剪数据
        if maxindex > 0:
            # 裁剪到maxindex+1列（包含maxindex）
            draw_data_cropped = draw_data[:, :maxindex+1]
            # 同时裁剪观测数据
            new_obs_list_cropped = [obs[:, :maxindex+1] for obs in new_obs_list]
        else:
            # 如果没有观测点，寻找数据的自然边界
            non_zero_cols = np.any(draw_data > 0, axis=0)
            if np.any(non_zero_cols):
                last_data_col = np.max(np.where(non_zero_cols)[0])
                crop_end = min(last_data_col + 3, draw_data.shape[1])
            else:
                crop_end = min(20, draw_data.shape[1])  # 默认20列
            
            draw_data_cropped = draw_data[:, :crop_end]
            new_obs_list_cropped = [obs[:, :crop_end] for obs in new_obs_list]
            
        # print(f"裁剪后数据形状: {draw_data_cropped.shape}")
        
        
        
        if  False:
            ''' 下面的代码是为了绘制 x - ini time , y - lead time '''
            
            # 🎯 步骤3：创建图形并绘制热图
            plt.figure(figsize=(15, 6))
            ax = sns.heatmap(draw_data_cropped/len(names), 
                             annot=True,
                             annot_kws={'size': 6}, 
                             fmt=".1f", 
                             cmap="YlGnBu", 
                             cbar_kws={'label': 'number of rapid intensification'}, 
                             linewidths=0.5,
                             vmin=0, 
                             vmax=20)
            
            # 🎯 步骤4：设置标签和格式
            ax.set_ylabel("Lead Time", fontsize=12)
            ax.set_xlabel("Initialization Time Minus First Initialization Time: h", fontsize=12)
            
            # 基于裁剪后的数据设置x轴标签
            dt = 12 if ini_time_mode=='00_12' else 24
            x_labels = [str(i * dt) for i in range(draw_data_cropped.shape[1])]
            ax.set_xticklabels(x_labels)
            
            plt.yticks(rotation=0)
            plt.xticks(rotation=90)
            
            # 确保y轴标签数量匹配
            ylabels = [str(i)+"d" for i in range(1, min(16, draw_data_cropped.shape[0] + 1))]
            ax.set_yticklabels(ylabels[:draw_data_cropped.shape[0]])
            
            ax.set_title(f'average total_RI Ensemble numbers of 8 typhoons east asia forecast_ini_mode:{ini_time_mode}')
            
            # 设置colorbar
            cbar = ax.collections[0].colorbar
            cbar.locator = ticker.MaxNLocator(integer=True)
            cbar.update_ticks()
            
            # 🎯 步骤5：在热图上绘制观测点
            if draw_obs_opt == True:
                for draw_data_obs in new_obs_list_cropped:
                    points = np.where(draw_data_obs == 1)
                    if len(points[1]) > 0:
                        y_coords = points[0] + 0.5
                        x_coords = points[1] + 0.15
                        ax.scatter(x_coords, y_coords, color='red', marker='o', s=10, zorder=5)
            
            # 保存和关闭
            if draw_opt == False:    
                plt.tight_layout()
                plt.savefig(f"{pic_savepath}/start_day_RIsum_81.png", dpi=900, bbox_inches='tight')
            else:
                plt.show()

            
            
        else:
            '''绘制 x-lead_time  y-init_time 图；最棘手的是红点逻辑全变 '''
            plt.figure(figsize=(15, 6))
            arr = draw_data_cropped.T
            ax = sns.heatmap(arr /len(names), 
                             annot=False,
                             annot_kws={'size': 6}, 
                             # 控制方格内数字显示的格式
                             fmt=".0f", 
                             cmap="YlGnBu", 
                             # format设置成整数，因为数组中的数值虽然是整数但类型是float，cbar刻度带小数点
                             cbar_kws={'label': 'number of rapid intensification','format':'%d'}, 
                             linewidths=0.5,
                             # 设置cbar刻度范围
                             vmin=0, 
                             vmax=20)
            # description
            ax.set_title(f'Mean number of Forecast RI members [{ini_time_mode} {obs_baseline}]')
            ax.set_xlabel('Lead Time')
            ax.set_ylabel('Relative to Unified RI Start Time (8 TCs)')
            
            # tick
            ax.set_xticklabels([f'{i+1}d' for i in range(arr.shape[1])])
            if ini_time_mode=='00_12':
                yticks      = [i+0.5 for i in range(arr.shape[0]) if (i*12)%48==0]
                yticklabels = [f"{(i*12)//24-14}d" for i in range(arr.shape[0]) if (i*12)%48==0]
                ax.set_yticks(yticks)
                ax.set_yticklabels(yticklabels, rotation = 0)
            else:
                yticks      = [i+0.5 for i in range(arr.shape[0]) if (i*24)%48==0]
                yticklabels = [f"{(i*24)//24-14}d" for i in range(arr.shape[0]) if (i*24)%48==0]
                ax.set_yticks(yticks)
                ax.set_yticklabels(yticklabels, rotation = 0)
            
            if draw_obs_opt == True:
                for draw_data_obs in new_obs_list_cropped:
                    points = np.where(draw_data_obs == 1)
                    if len(points[1]) > 0:
                        x_coords = points[0] + 0.5
                        y_coords = points[1] + 0.5
                        ax.scatter(x_coords, y_coords, color='red', marker='o', s=10, zorder=5)
            ax.invert_yaxis()
            
            # 保存和关闭
            plt.tight_layout()
            if draw_opt==True:
                plt.show()
            else:
                plt.savefig(os.path.join(pic_savepath,'start_day_RIsum_81.png'), dpi=900, bbox_inches='tight')
                plt.close()
            
            
